class Module:
    def __init__(self):
        self.__modules = []

    def forward(self, x):
        return x

    def __call__(self, x):
        return self.forward(x)

    def state_dict(self):
        weights = []
        for module in self.__modules:
            weights.append(module.get_weights())
        return weights

    def load_state_dict(self, data):
        for module in self.__modules:
            w_cnt = module.set_weights(data)
            del data[0:w_cnt]
        assert len(data) == 0, 'Keys mismatch.'

    def get_weights(self):
        pass

    def set_weights(self, data):
        return 0

    def __setattr__(self, key, value):
        if isinstance(value, Module):
            self.__modules.append(value)
        self.__dict__[key] = value
